#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
更新指定conda环境中的包，基于pip requirements文件
"""

import os
import sys
import subprocess
import argparse
from pathlib import Path

def list_conda_environments():
    """列出所有conda环境"""
    result = subprocess.run(["conda", "env", "list"], capture_output=True, text=True)
    
    if result.returncode != 0:
        print(f"获取conda环境列表失败: {result.stderr}")
        return []
    
    envs = []
    for line in result.stdout.split('\n'):
        if line.startswith('#') or not line.strip():
            continue
        
        parts = line.split()
        if len(parts) >= 1:
            envs.append(parts[0])
    
    return envs

def list_available_requirements():
    """列出可用的pip requirements文件"""
    env_dir = Path("conda_environments")
    if not env_dir.exists():
        print(f"错误: 找不到环境目录 {env_dir}")
        return []
    
    req_files = list(env_dir.glob("*_pip_requirements.txt"))
    return [f.name for f in req_files]

def update_environment(env_name, req_file=None):
    """更新指定conda环境中的包"""
    # 检查环境是否存在
    envs = list_conda_environments()
    if env_name not in envs:
        print(f"错误: 环境 {env_name} 不存在")
        return False
    
    # 如果没有指定requirements文件，尝试使用与环境名匹配的文件
    if req_file is None:
        req_path = Path("conda_environments") / f"{env_name}_pip_requirements.txt"
        if not req_path.exists():
            print(f"错误: 未找到与环境 {env_name} 匹配的pip requirements文件")
            print(f"请先使用 'make export-pip' 导出环境的pip requirements")
            return False
    else:
        req_path = Path(req_file)
        if not req_path.exists():
            print(f"错误: requirements文件 {req_file} 不存在")
            return False
    
    print(f"正在使用文件 {req_path} 更新conda环境 {env_name} 中的包...")
    
    # 使用conda run命令在指定环境中运行pip install
    process = subprocess.run(["conda", "run", "-n", env_name, "pip", "install", "-r", str(req_path), "--upgrade"])
    
    if process.returncode != 0:
        print(f"环境更新失败，返回码: {process.returncode}")
        return False
    
    print(f"环境 {env_name} 更新成功!")
    return True

def main():
    parser = argparse.ArgumentParser(description="更新指定conda环境中的包，基于pip requirements文件")
    parser.add_argument("--env-name", type=str, help="要更新的conda环境名称")
    parser.add_argument("--req-file", type=str, help="pip requirements文件路径")
    parser.add_argument("--list-envs", action="store_true", help="列出所有conda环境")
    parser.add_argument("--list-reqs", action="store_true", help="列出所有可用的pip requirements文件")
    
    args = parser.parse_args()
    
    if args.list_envs:
        print("可用的conda环境:")
        for env in list_conda_environments():
            print(f"- {env}")
        return
    
    if args.list_reqs:
        print("可用的pip requirements文件:")
        for req in list_available_requirements():
            print(f"- {req}")
        return
    
    if not args.env_name:
        # 如果没有指定环境名称，列出可用选项并让用户选择
        available_envs = list_conda_environments()
        if not available_envs:
            print("没有找到任何conda环境")
            return
        
        print("请选择要更新的conda环境:")
        for i, env in enumerate(available_envs):
            print(f"{i+1}. {env}")
        
        try:
            choice = int(input("请输入选项编号: "))
            if 1 <= choice <= len(available_envs):
                env_name = available_envs[choice-1]
            else:
                print("无效的选择")
                return
        except ValueError:
            print("请输入有效的数字")
            return
    else:
        env_name = args.env_name
    
    # 如果没有指定requirements文件，但用户想要选择一个
    if not args.req_file and input("是否要指定requirements文件? (y/n): ").lower() == 'y':
        available_reqs = list_available_requirements()
        if not available_reqs:
            print("没有找到任何pip requirements文件")
            print("请先使用 'make export-pip' 导出环境的pip requirements")
            return
        
        print("请选择要使用的pip requirements文件:")
        for i, req in enumerate(available_reqs):
            print(f"{i+1}. {req}")
        
        try:
            choice = int(input("请输入选项编号: "))
            if 1 <= choice <= len(available_reqs):
                req_file = str(Path("conda_environments") / available_reqs[choice-1])
            else:
                print("无效的选择")
                return
        except ValueError:
            print("请输入有效的数字")
            return
    else:
        req_file = args.req_file
    
    update_environment(env_name, req_file)

if __name__ == "__main__":
    main() 